package hebClustering.clusteringAlgorithms;

import hebClustering.Cluster;
import hebClustering.ClusterSet;
import hebClustering.vectorSpace.IVector;
import hebClustering.vectorSpace.distances.IDistance;

import java.util.*;

import org.jgrapht.*;
import org.jgrapht.alg.ConnectivityInspector;
import org.jgrapht.alg.KruskalMinimumSpanningTree;

import org.jgrapht.graph.*;

/**
 * @author Roman Tzipron
 * The MST algorithm turns the metric space into a complete weighted graph as follows: the vectors are regarded as vertexes and each vertex connected to every other vertex through an edge which weight is the 
 * 	metric distance between the two vertexes(vectors) it connects.
 * A minimum spanning tree of the resulting graph is computed.
 * After that an iterative process deletes edges which meet a certain criterion. 
 * Each removal turns one connected component into two(since it is a tree). 
 * After the process finishes deleting the edges we are left with a forest of trees each of which represents a cluster.
 */

public class MSTClustering implements IClusteringAlgorithm {

	private int K;
	IDistance distance;

	public MSTClustering(IDistance d,int K) {
		distance = d;
		this.K = K;
	}

	@Override
	public ClusterSet cluster(Collection<IVector> dataSet) {
		UndirectedWeightedSubgraph<IVector, DefaultWeightedEdge> distanceMST =
			computeDistanceMST(generateDistanceGraph(dataSet));
		
		
		for(int i=1; i<K; i++){
			distanceMST.removeEdge(bestRemovalEdge(distanceMST));
		}
		
		return clustersFromTree(distanceMST);
	}


	private ClusterSet clustersFromTree(
			UndirectedWeightedSubgraph<IVector, DefaultWeightedEdge> distanceMST) {
		ClusterSet clusterSet = new ClusterSet();
		
		ConnectivityInspector<IVector, DefaultWeightedEdge> cInspector = new ConnectivityInspector<IVector, DefaultWeightedEdge>(distanceMST);
		List< Set<IVector> > clustersList = cInspector.connectedSets();
		
		for(Set<IVector> s : clustersList){
			Cluster cluster = new Cluster();
			clusterSet.add(cluster);
			cluster.addAll(s);
		}
		
		return clusterSet;
	}

	private DefaultWeightedEdge bestRemovalEdge(
			UndirectedWeightedSubgraph<IVector, DefaultWeightedEdge> distanceMST) {
		
		DefaultWeightedEdge bestEdge = null;
		double bestSoFar = 0;
		
		for(DefaultWeightedEdge e : distanceMST.edgeSet()){

			UndirectedWeightedSubgraph<IVector, DefaultWeightedEdge> graphCopy = 
				new UndirectedWeightedSubgraph<IVector, DefaultWeightedEdge>(distanceMST, distanceMST.vertexSet(), distanceMST.edgeSet());
			
			graphCopy.removeEdge(e);
			
			WeightedGraph<IVector, DefaultWeightedEdge> tree1,tree2;
			tree1 = treeByVertex(distanceMST.getEdgeSource(e),graphCopy);
			tree2 = treeByVertex(distanceMST.getEdgeTarget(e),graphCopy);
			
			double edgeScore = distanceMST.getEdgeWeight(e)*Math.log(treeWeight(tree1)*treeWeight(tree2));
				
			if(bestSoFar<edgeScore){
				bestSoFar = edgeScore;
				bestEdge = e;
			}
		}
		
		return bestEdge;
	}

	private WeightedGraph<IVector, DefaultWeightedEdge> treeByVertex(IVector v,UndirectedWeightedSubgraph<IVector, DefaultWeightedEdge> graph) {
		
		ConnectivityInspector<IVector, DefaultWeightedEdge> cInspector = new ConnectivityInspector<IVector, DefaultWeightedEdge>(graph);
		Set<IVector> vectorSet = cInspector.connectedSetOf(v);
		return new UndirectedWeightedSubgraph<IVector, DefaultWeightedEdge>(graph, vectorSet, spanningEdges(graph, vectorSet));
	}

	private double treeWeight(WeightedGraph<IVector, DefaultWeightedEdge> tree) {
		double sum = 0;
		
		for(DefaultWeightedEdge e : tree.edgeSet()){
			sum+=tree.getEdgeWeight(e);
		}
		
		return sum;
	}

	private Set<DefaultWeightedEdge> spanningEdges(UndirectedSubgraph<IVector,DefaultWeightedEdge> graph,Set<IVector> vertexs) {
		Set<DefaultWeightedEdge> span = new HashSet<DefaultWeightedEdge>();

		for(IVector v : vertexs){
			span.addAll(graph.edgesOf(v));
		}

		return span;
	}

	private UndirectedWeightedSubgraph<IVector, DefaultWeightedEdge> computeDistanceMST(
			WeightedGraph<IVector, DefaultWeightedEdge> distanceGraph) {

		KruskalMinimumSpanningTree<IVector,DefaultWeightedEdge> kruskalMST = new KruskalMinimumSpanningTree<IVector, DefaultWeightedEdge>(distanceGraph);
		return new UndirectedWeightedSubgraph<IVector,DefaultWeightedEdge>(distanceGraph,distanceGraph.vertexSet(),kruskalMST.getEdgeSet());
	}

	private WeightedGraph<IVector, DefaultWeightedEdge> generateDistanceGraph(Collection<IVector> dataSet) {

		WeightedGraph<IVector,DefaultWeightedEdge> graph = new SimpleWeightedGraph<IVector, DefaultWeightedEdge>(DefaultWeightedEdge.class);

		for(IVector v1 : dataSet){
			for(IVector v2 : dataSet){
				if(!graph.containsEdge(v1, v2) && v1!=v2){
					graph.addVertex(v1);
					graph.addVertex(v2);
					DefaultWeightedEdge edge = graph.addEdge(v1, v2);
					graph.setEdgeWeight(edge, distance.calc(v1, v2));

				}
			}
		}

		return graph;
	}

	public String toString(){
		return "MST";
	}
}
